from ultralytics import YOLO
from PIL import Image, ImageDraw
import torch


class Detector:
    def __init__(self, model=None):
        if model:
            import torch.serialization
            from ultralytics.nn.tasks import ClassificationModel
            with torch.serialization.safe_globals([ClassificationModel]):
                self.model = YOLO(model)
            device = "cuda" if torch.cuda.is_available() else "cpu"
            self.model.to(device)

    def predict(self, image):
        """预测一帧"""
        res = self.model.predict(image)
        labels = res[0].names
        probs = res[0].probs.cpu().numpy()
        best_result = probs.top1
        confidence = probs.top1conf
        return best_result, confidence, labels[best_result]

    def predict_image(self, image):
        """test"""
        best_result, confidence, class_name = self.predict(image)
        image = image.convert("RGB")
        draw = ImageDraw.Draw(image)
        confidence = confidence * 100
        text = "{:<5} {:>.2f}%".format(class_name, confidence)
        draw.text((50, 50), text, fill=(255, 0, 0, 1))
        return best_result, confidence, image

    def info(self):
        """test"""
        print("=" * 10, "模型参数", "=" * 10)
        print(self.model)
        if hasattr(self.model, "names"):
            print("Classes:", self.model.names)
        else:
            print("No class names attribute found.")
        print("=" * 30)
